#!/usr/bin/env bash
# Usage examples for tools/extract_attention.py
# Ensure you have installed dependencies:
#   pip install requests
# and installed the package (if needed):
#   pip install -e .
#
# Variables to set before running:
#   CKPT - path to the model checkpoint (.pt)
#   OUTDIR - path where h5 outputs will be written
#   LINKS_JSON - path to co3d links.json (contains per-category zip urls)
#   CO3D_ANNOT - path to CO3D annotation directory containing <category>_train.jgz files
#   WORK_DIR - optional working directory for downloads and extraction (default: /tmp/co3d_work)
#
# Example 1: single-category quick test (download only 'apple', process 2 samples)
# Replace the paths with your actual values.
# This will download the apple zips, extract to WORK_DIR, run model inference for max 2 samples, write HDF5 and then cleanup extracted files.

CKPT=${CKPT:-/path/to/checkpoint.pt}
OUTDIR=${OUTDIR:-/path/to/out_attn}
LINKS_JSON=${LINKS_JSON:-/data1/lqf/co3d/co3d/links.json}
CO3D_ANNOT=${CO3D_ANNOT:-/path/to/CO3D_ANNOT}
WORK_DIR=${WORK_DIR:-/tmp/co3d_work}

mkdir -p "$OUTDIR"
mkdir -p "$WORK_DIR"

python tools/extract_attention.py \
  --ckpt "$CKPT" \
  --out_dir "$OUTDIR" \
  --links_json "$LINKS_JSON" \
  --CO3D_ANNOT_DIR "$CO3D_ANNOT" \
  --work_dir "$WORK_DIR" \
  --category apple \
  --max_samples 2 \
  --batch_size 1 \
  --num_workers 2 \
  --device cpu \
  --single_h5 \
  --head_avg

# Example 2: iterate all categories in links.json (be careful: will download each category sequentially)
# This will process all categories listed in the 'full' section of links.json. Use --max_samples to limit per-category processing.
# Note: this can take a long time and requires network and disk space for temporary extraction per category.

# python tools/extract_attention.py \
#   --ckpt "$CKPT" \
#   --out_dir "$OUTDIR" \
#   --links_json "$LINKS_JSON" \
#   --CO3D_ANNOT_DIR "$CO3D_ANNOT" \
#   --work_dir "$WORK_DIR" \
#   --max_samples 100 \
#   --batch_size 1 \
#   --num_workers 4 \
#   --device cuda \
#   --single_h5 \
#   --head_avg

# Troubleshooting tips:
# - If the script fails to find a model class, open tools/extract_attention.py and change the explicit import to your model class (e.g. from vggt.models.vggt import VGGT)
# - Ensure CO3D annotation files are present in CO3D_ANNOT (e.g. apple_train.jgz, apple_test.jgz)
# - If downloads keep failing, try downloading manually into WORK_DIR and re-run with the same WORK_DIR to use cached zips

# End of script
